package cn.monkey.state.scheduler;

import cn.monkey.state.bean.Countable;
import cn.monkey.state.bean.Refreshable;
import cn.monkey.state.core.FetchStateGroup;
import cn.monkey.state.core.StateGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public abstract class AbstractSchedulerManager implements SchedulerManager, Countable, Closeable, Refreshable {

    private static final Logger log = LoggerFactory.getLogger(AbstractSchedulerManager.class);

    protected final StateGroupSchedulerFactory stateGroupSchedulerFactory;

    protected final EventPublishSchedulerFactory eventPublishSchedulerFactory;

    protected final EventPublishScheduler[] eventPublishSchedulers;

    protected final SchedulerManagerConfig schedulerManagerConfig;

    protected final AtomicBoolean isRunning;
    protected final AtomicLong stateGroupSchedulerIdCounter = new AtomicLong(0);

    protected volatile Map<Long, StateGroupScheduler> stateGroupSchedulerMap;


    public AbstractSchedulerManager(StateGroupSchedulerFactory stateGroupSchedulerFactory,
                                    EventPublishSchedulerFactory eventPublishSchedulerFactory,
                                    SchedulerManagerConfig managerConfig) {
        this.stateGroupSchedulerFactory = stateGroupSchedulerFactory;
        this.eventPublishSchedulerFactory = eventPublishSchedulerFactory;
        this.schedulerManagerConfig = managerConfig;
        this.eventPublishSchedulers = this.initEventPublishSchedulers();
        this.stateGroupSchedulerMap = this.createStateGroupSchedulerMap();
        this.isRunning = new AtomicBoolean(true);
    }

    protected ConcurrentHashMap<Long, StateGroupScheduler> createStateGroupSchedulerMap() {
        int stateGroupSchedulerCoreSize = this.schedulerManagerConfig.getStateGroupSchedulerCoreSize();
        ConcurrentHashMap<Long, StateGroupScheduler> stateGroupSchedulerMap = new ConcurrentHashMap<>();
        if (stateGroupSchedulerCoreSize > 0) {
            for (int i = 0; i < stateGroupSchedulerCoreSize; i++) {
                StateGroupScheduler stateGroupScheduler = this.stateGroupSchedulerFactory.create(i);
                stateGroupScheduler.start();
                stateGroupSchedulerMap.put(stateGroupScheduler.id(), stateGroupScheduler);
            }
        }
        return stateGroupSchedulerMap;
    }

    protected EventPublishScheduler[] initEventPublishSchedulers() {
        int eventPublisherSchedulerSize = this.schedulerManagerConfig.getEventPublisherSchedulerSize();
        if (eventPublisherSchedulerSize != 1 && Integer.bitCount(eventPublisherSchedulerSize) != 1) {
            throw new IllegalArgumentException("event publisher size must be a power of 2 or equals to 1");
        }
        EventPublishScheduler[] eventPublishSchedulers = new EventPublishScheduler[eventPublisherSchedulerSize];
        for (int i = 0; i < eventPublisherSchedulerSize; i++) {
            eventPublishSchedulers[i] = this.eventPublishSchedulerFactory.create(i);
            eventPublishSchedulers[i].start();
        }
        return eventPublishSchedulers;
    }

    protected final EventPublishScheduler findEventPublisherScheduler(String groupId) {
        int length = this.eventPublishSchedulers.length;
        return this.eventPublishSchedulers[groupId.hashCode() & (length - 1)];
    }

    @Override
    public boolean isEmpty() {
        return this.size() == 0;
    }

    @Override
    public int size() {
        return this.stateGroupSchedulerMap.size();
    }

    @Override
    public boolean isFull() {
        return this.size() >= this.schedulerManagerConfig.getStateGroupSchedulerSize();
    }

    protected void onFetchStateGroup(FetchStateGroup fetchStateGroup, Object event) {
        StateGroup stateGroup = fetchStateGroup.stateGroup();
        if (!fetchStateGroup.isNew()) {
            stateGroup.addEvent(event);
            return;
        }
        for (StateGroupScheduler scheduler : this.stateGroupSchedulerMap.values()) {
            if (scheduler.isFull()) {
                continue;
            }
            if (!scheduler.isStarted()) {
                continue;
            }
            if (scheduler.tryAddStateGroup(stateGroup)) {
                stateGroup.addEvent(event);
                return;
            }
        }
        if (this.isFull()) {
            log.error("can not add any stateGroupScheduler");
            throw new IllegalArgumentException("can not add any stateGroupScheduler");
        }
        long id = this.stateGroupSchedulerIdCounter.getAndIncrement();
        if (id >= this.schedulerManagerConfig.getStateGroupSchedulerSize()) {
            log.error("can not create any stateGroupScheduler");
            return;
        }
        this.stateGroupSchedulerMap.computeIfAbsent(id, (k) -> {
            StateGroupScheduler stateGroupScheduler = this.stateGroupSchedulerFactory.create(k);
            stateGroupScheduler.start();
            stateGroupScheduler.tryAddStateGroup(stateGroup);
            return stateGroupScheduler;
        });
    }


    @Override
    public void refresh() {
        int coreSize = this.schedulerManagerConfig.getStateGroupSchedulerCoreSize();
        int currentSize = this.size();
        if (currentSize <= coreSize) {
            return;
        }
        int removableCount = currentSize - coreSize;
        int removedCount = 0;
        Iterator<Map.Entry<Long, StateGroupScheduler>> iterator = this.stateGroupSchedulerMap.entrySet().iterator();
        while (iterator.hasNext() && removedCount < removableCount) {
            Map.Entry<Long, StateGroupScheduler> entry = iterator.next();
            StateGroupScheduler scheduler = entry.getValue();
            if (scheduler != null && scheduler.isEmpty()) {
                scheduler.stop();
                iterator.remove();
                removedCount++;
            }
        }
    }

    @Override
    public void close() throws IOException {
        this.isRunning.set(false);
        Collection<StateGroupScheduler> values = this.stateGroupSchedulerMap.values();
        for (StateGroupScheduler scheduler : values) {
            scheduler.stop();
        }
        this.stateGroupSchedulerMap.clear();
        for (EventPublishScheduler scheduler : this.eventPublishSchedulers) {
            scheduler.stop();
        }
        Arrays.fill(this.eventPublishSchedulers, null);
    }
}
